'''
Created on May 25, 2009

@author: user
'''
from cliques.potentialBasedSearch import potentialBasedSearch
import cliques.mdp.mdpSolver
import logging
class mdpBasedSearch(potentialBasedSearch):
    '''
    classdocs
    '''

    def __init__(self,mdp_solver = cliques.mdp.mdpSolver.cliqueMdpSolver,name = "mdpBasedSearch",lookahead_depth=2,states_per_node=10,max_states=100000, \
                 noise = 0,constant_edge_prob=0):
        '''
        Constructor
        '''
        potentialBasedSearch.__init__(self)
        self.solverCreator = mdp_solver
        self.name = name
        self.lookahead_depth=lookahead_depth
        self.states_per_node = states_per_node
        self.max_states = max_states
        self.noise = noise
        self.p=constant_edge_prob

    def setup(self, unknown_graph, k, start_node):
        potentialBasedSearch.setup(self, unknown_graph, k, start_node)

        # Calculate edge prob
        nodes = float(len(unknown_graph.nodes()))
        edges = float(len(unknown_graph.edges()))
          
        self.p =  edges/(nodes*(nodes-1)) # SHOULD DIVIDE BY 2?


        # Create an MDP solver
        self.solver = self.solverCreator(self.k, self.p, nodes)
        if "noise" in dir(self.solver):
            self.solver.noise = self.noise
        if "max_states_per_action" in dir(self.solver):
            self.solver.max_states_per_action=self.states_per_node
        if "max_states" in dir(self.solver):
            self.solver.max_states = self.max_states
        if "real_graph" in dir(self.solver):
            self.solver.real_graph = unknown_graph
        if "num_of_trials" in dir(self.solver):
            self.solver.num_of_trials=self.states_per_node
                                                
        if "set_search" in dir(self.solver): 
            self.solver.set_search(self)            
        if "set_lookahead_depth" in dir(self.solver):
            self.solver.set_lookahead_depth(self.lookahead_depth)
        if "set_real_graph" in dir(self.solver):
            self.solver.set_real_graph(unknown_graph)
            
    def choose_node(self):
        if self.iteration<2: # First 2 iterations are random
            return self.generated[0]
        
        # create initial mdp state
        (v,policy) = self.solver.run(self.ckg, self.generated)
        self.state_0 = self.solver.initial
        self.current_v = v
        self.current_policy = policy
        #self.state_graph = self.solver.state_graph
        logging.debug("Chosen node %s with expected cost of %f" % (policy[self.state_0],v[self.state_0]))
        if policy[self.state_0]==None:
            return self.default_choose_node()
        return policy[self.state_0]
    
    def default_choose_node(self):
        ''' When all states look the same (no potential cliques exist) - use known degree '''        
        max_degree = 0
        
        best_node = None
        for node in self.generated:
            degree=self.ckg.order(node)
            if max_degree<degree:
                max_degree=degree
                best_node=node
        return best_node
        
    def __str__(self):
        return self.name
    
                